import torch
from torch import nn
import torch.nn.functional as F
import numpy as np

import torchvision.models._utils as _utils
from functools import partial



def bbox_areas(bboxes, keep_axis=False):
    x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]
    areas = (y_max - y_min + 1) * (x_max - x_min + 1)
    if keep_axis:
        return areas[:, None]
    return areas

def gaussian_2d(shape, sigma_x=1, sigma_y=1):

    m, n = [(ss - 1.) / 2. for ss in shape]
    y, x = np.ogrid[-m:m + 1, -n:n + 1]

    h = np.exp(-(x * x / (2 * sigma_x * sigma_x) + y * y / (2 * sigma_y * sigma_y)))
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    return h
def draw_truncate_gaussian(heatmap, center, h_radius, w_radius, k=1):
    h, w = 2 * h_radius + 1, 2 * w_radius + 1
    sigma_x = w / 6
    sigma_y = h / 6
    gaussian = gaussian_2d((h, w), sigma_x=sigma_x, sigma_y=sigma_y)
    gaussian = heatmap.new_tensor(gaussian)   #即 gaussian

    x, y = int(center[0]), int(center[1])
    print('before:',heatmap)
    height, width = heatmap.shape[0:2]

    left, right = min(x, w_radius), min(width - x, w_radius + 1)
    top, bottom = min(y, h_radius), min(height - y, h_radius + 1)

    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
    masked_gaussian = gaussian[h_radius - top:h_radius + bottom,
                      w_radius - left:w_radius + right]
    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:
        torch.max(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
    print('after:',heatmap)
    return heatmap

def target_single_image(gt_boxes, gt_labels):
    output_h = int(40//4)
    output_w =  int(40//4)
    heatmap = gt_boxes.new_zeros((1, output_h, output_w))
    fake_heatmap = gt_boxes.new_zeros((output_h, output_w))
    box_target = gt_boxes.new_zeros((4, output_h, output_w))
    reg_weight = gt_boxes.new_zeros((1, output_h, output_w))

    boxes_areas_log = bbox_areas(gt_boxes).log()
    boxes_area_topk_log, boxes_ind = torch.topk(boxes_areas_log, boxes_areas_log.size(0))
    gt_boxes = gt_boxes[boxes_ind]
    gt_labels = gt_labels[boxes_ind]

    feat_gt_boxes = gt_boxes / 4
    feat_gt_boxes[:, [0, 2]] = torch.clamp(feat_gt_boxes[:, [0, 2]], min=0,
                                           max=output_w - 1)
    feat_gt_boxes[:, [1, 3]] = torch.clamp(feat_gt_boxes[:, [1, 3]], min=0,
                                           max=output_h - 1)
    feat_hs, feat_ws = (feat_gt_boxes[:, 3] - feat_gt_boxes[:, 1],
                        feat_gt_boxes[:, 2] - feat_gt_boxes[:, 0])
    
    ct_ints = (torch.stack([(gt_boxes[:, 0] + gt_boxes[:, 2]) / 2,
                            (gt_boxes[:, 1] + gt_boxes[:, 3]) / 2],
                           dim=1) / 4).to(torch.int)
    h_radiuses_alpha = (feat_hs / 2. * 0.54).int()
    w_radiuses_alpha = (feat_ws / 2. * 0.54).int()

    for k in range(boxes_ind.shape[0]):
        cls_id = gt_labels[k]

        # print(cls_id)

        fake_heatmap = fake_heatmap.zero_()
        draw_truncate_gaussian(fake_heatmap, ct_ints[k],
                h_radiuses_alpha[k].item(), w_radiuses_alpha[k].item())
        #print(heatmap.shape,cls_id,fake_heatmap.shape)
        print('fake:',fake_heatmap)
        heatmap[cls_id-1] = torch.max(heatmap[cls_id-1], fake_heatmap)
        
        box_target_inds = fake_heatmap > 0
        
        box_target[:, box_target_inds] = gt_boxes[k][:, None] #
        print('box_target_inds:',box_target_inds)
        print(box_target)

        local_heatmap = fake_heatmap[box_target_inds]
        ct_div = local_heatmap.sum()
        local_heatmap *= boxes_area_topk_log[k]
        reg_weight[0, box_target_inds] = local_heatmap / ct_div
        print(reg_weight)

    return heatmap, box_target, reg_weight

gt_boxes = torch.tensor([[3.0,4.0,28.0,24.0],[15.0,21.0,36.0,50.0]])
gt_labels = torch.tensor([[0],[0]])
heatmap, box_target, reg_weight = target_single_image(gt_boxes,gt_labels)